#include "Scene.h"

#include "utils/D3DUtils.h"
#include "utils/DxErrorChecker.h"
#include "utils/MathUtils.h"

#include "Globals.h"

using namespace DirectX;

void Scene::init() {
    // Directional lights.
    mLight.mAmbient = XMFLOAT4A(0.3f, 0.3f, 0.3f, 1.0f);
    mLight.mDiffuse = XMFLOAT4A(1.0f, 1.0f, 1.0f, 1.0f);
    mLight.mSpecular = XMFLOAT4A(0.8f, 0.8f, 0.7f, 1.0f);
    mLight.mDir = XMFLOAT3A(0.0f, -1.0f, 0.0f);

    // Initialize material
    mQuadMaterial.mAmbient = XMFLOAT4A(1.0f, 1.0f, 1.0f, 1.0f);
    mQuadMaterial.mDiffuse = XMFLOAT4A(1.0f, 1.0f, 1.0f, 1.0f);
    mQuadMaterial.mSpecular = XMFLOAT4A(0.0f, 0.0f, 0.0f, 64.0f);
    mQuadMaterial.mReflect = XMFLOAT4A(0.0f, 0.0f, 0.0f, 1.0f);

    // Initialize constant buffers
    assert(Globals::gD3D.mDevice);
    ID3D11Device& device = *Globals::gD3D.mDevice;

    const float scalingFactor = 1.0f;
    const XMMATRIX scaling = XMMatrixScaling(scalingFactor, scalingFactor, scalingFactor);
    XMStoreFloat4x4(&mTexW, scaling);

    const XMMATRIX texW = XMLoadFloat4x4(&mTexW);
    const XMMATRIX texWTranspose = XMMatrixTranspose(texW);
    XMFLOAT4X4 textureW;
    XMStoreFloat4x4(&textureW, texWTranspose);
    D3D11_SUBRESOURCE_DATA initData;
    initData.pSysMem = &textureW;
    initData.SysMemPitch = 0;
    initData.SysMemSlicePitch = 0;

    HRESULT result = createBuffer(D3D11_USAGE_IMMUTABLE,
                                  D3D11_BIND_CONSTANT_BUFFER,
                                  0,
                                  0,
                                  0,
                                  sizeof(mVSImmutable.mData),
                                  &initData,
                                  mVSImmutable.mBuffer);

    result = createBuffer(D3D11_USAGE_DYNAMIC,
                          D3D11_BIND_CONSTANT_BUFFER,
                          D3D11_CPU_ACCESS_WRITE,
                          0,
                          0,
                          sizeof(mVSPerFrameBuffer.mData),
                          nullptr,
                          mVSPerFrameBuffer.mBuffer);

    DxErrorChecker(result);

    result = createBuffer(D3D11_USAGE_DYNAMIC,
                          D3D11_BIND_CONSTANT_BUFFER,
                          D3D11_CPU_ACCESS_WRITE,
                          0,
                          0,
                          sizeof(mPsPerFrameBuffer.mData),
                          nullptr,
                          mPsPerFrameBuffer.mBuffer);

    DxErrorChecker(result);

    // World matrix
    set(0.0f, 0.0f, 0.0f, mQuadCenterPosW);
    set(1.0f, 0.0f, 0.0f, mQuadVelW);

    //
    // Init rendering and simulation frames
    //
    FrameInfo* forRendering = Globals::gFrameBuffer.mForRendering;
    FrameInfo* forSimulation = Globals::gFrameBuffer.mForSimulation;

    set(2.0f, 0.0f, 2.0f, forRendering->mEntitiesData[0].mCenterPosW);
    set(1.0f, 0.0f, 1.0f, forRendering->mEntitiesData[0].mVelW);
    set(-2.0f, 0.0f, 2.0f, forRendering->mEntitiesData[1].mCenterPosW);
    set(-1.0f, 0.0f, 1.0f, forRendering->mEntitiesData[1].mVelW);

    set(2.0f, 0.0f, 2.0f, forSimulation->mEntitiesData[0].mCenterPosW);
    set(1.0f, 0.0f, 1.0f, forSimulation->mEntitiesData[0].mVelW);
    set(-2.0f, 0.0f, 2.0f, forSimulation->mEntitiesData[1].mCenterPosW);
    set(-1.0f, 0.0f, 1.0f, forSimulation->mEntitiesData[1].mVelW);
}

void Scene::draw() {
    assert(Globals::gD3D.mContext);
    ID3D11DeviceContext& context = *Globals::gD3D.mContext;
    const uint32_t maxConstantBufferSlots = D3D11_COMMONSHADER_CONSTANT_BUFFER_API_SLOT_COUNT ;
    const uint32_t maxResourceSlots = D3D11_COMMONSHADER_INPUT_RESOURCE_SLOT_COUNT;
    const uint32_t maxSamplerSlots = D3D11_COMMONSHADER_SAMPLER_SLOT_COUNT;
    
    //
    // Input Assembler Stage
    //

    // Input Layout 
    ID3D11InputLayout * const inputLayout = Globals::gShaders.mIL;
    context.IASetInputLayout(inputLayout);

    // Primitive Topology
    context.IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);

    // Vertex Buffer
    ID3D11Buffer* buffers[2] = 
        { 
            Globals::gGeomBuffers.mBufferInfo->mVertexBuffer,
            Globals::gGeomBuffers.mInstanceBuffer.mBuffer.mBuffer
        };
    uint32_t strides[2] = { sizeof(Vertex), sizeof(KinematicInfo) };
    uint32_t offsets[2] = { 0, 0 };
    context.IASetVertexBuffers(0, 2, buffers, strides, offsets);

    // Index Buffer
    ID3D11Buffer* indexBuffer = Globals::gGeomBuffers.mBufferInfo->mIndexBuffer;
    context.IASetIndexBuffer(indexBuffer, DXGI_FORMAT_R32_UINT, 0);

    //
    // Vertex Shader Stage
    //

    // Shader
    ID3D11VertexShader * const vertexShader = Globals::gShaders.mVS;
    context.VSSetShader(vertexShader, nullptr, 0);

    // Per Frame Constant Buffer
    const XMMATRIX viewProj = Globals::gCamera.viewProjMatrix();
    const XMMATRIX viewProjTranspose = XMMatrixTranspose(viewProj);
    XMStoreFloat4x4A(&mVSPerFrameBuffer.mData.mViewProj, viewProjTranspose);
    copyData(&mVSPerFrameBuffer.mData, sizeof(mVSPerFrameBuffer.mData), *mVSPerFrameBuffer.mBuffer); 

    // Set Constant Buffers
    ID3D11Buffer* vertexShaderBuffers[maxConstantBufferSlots] = { nullptr };
    vertexShaderBuffers[0] = mVSPerFrameBuffer.mBuffer;
    vertexShaderBuffers[1] = mVSImmutable.mBuffer;
    context.VSSetConstantBuffers(0, maxConstantBufferSlots, vertexShaderBuffers);        

    //
    // Pixel Shader Stage
    //

    // Shader
    ID3D11PixelShader* const pixelShader = Globals::gShaders.mPS;
    context.PSSetShader(pixelShader, nullptr, 0);

    // Per Frame Constant Buffer
    memcpy(&mPsPerFrameBuffer.mData.mLight, &mLight, sizeof(mLight));
    mPsPerFrameBuffer.mData.mEyePosW = Globals::gCamera.mPosW;
    mPsPerFrameBuffer.mData.mMaterial = mQuadMaterial;
    copyData(&mPsPerFrameBuffer.mData, sizeof(mPsPerFrameBuffer.mData), *mPsPerFrameBuffer.mBuffer);

    // Set constant buffers
    ID3D11Buffer* pixelShaderBuffers[maxConstantBufferSlots] = { nullptr };
    pixelShaderBuffers[0] = mPsPerFrameBuffer.mBuffer;
    context.PSSetConstantBuffers(0, maxConstantBufferSlots, pixelShaderBuffers);

    // Resources
    ID3D11ShaderResourceView* pixelShaderResources[maxResourceSlots] = {nullptr};
    pixelShaderResources[0] = Globals::gResources.mQuadTextureSRV;
    context.PSSetShaderResources(0, maxResourceSlots, pixelShaderResources);

    // Sampler state
    ID3D11SamplerState* samplerStates[maxSamplerSlots] = {nullptr};
    samplerStates[0] = Globals::gPipeline.mAnisotropicSS;
    context.PSSetSamplers(0, maxSamplerSlots, samplerStates);

    //
    // Draw
    // 

    updateInstanceBuffers();

    // Blend factors to be used in blending
    //float blendFactor[] = {0.0f, 0.0f, 0.0f, 0.0f};

    const uint32_t baseVertexLocation = Globals::gGeomBuffers.mBufferInfo->mBaseVertexLocation;
    const uint32_t startIndexLocation = Globals::gGeomBuffers.mBufferInfo->mStartIndexLocation;
    const uint32_t indexCount = Globals::gGeomBuffers.mBufferInfo->mIndexCount;
    //context.OMSetBlendState(Globals::gPipeline.mTransparentBS, blendFactor, 0xffffffff);
    context.DrawIndexedInstanced(indexCount, 
                                 Globals::gGeomBuffers.mInstanceBuffer.mNumInstances, 
                                 startIndexLocation, 
                                 baseVertexLocation,
                                 0);

    // Restore default blend state
    //context.OMSetBlendState(0, blendFactor, 0xffffffff);
}

void Scene::destroy() {
    mVSImmutable.mBuffer->Release();
    mVSPerFrameBuffer.mBuffer->Release();
    mPsPerFrameBuffer.mBuffer->Release();
}

void Scene::updateInstanceBuffers() {
    KinematicInfo* &instanceData = Globals::gGeomBuffers.mInstanceBuffer.mBuffer.mData;
    const uint32_t numInstances = Globals::gGeomBuffers.mInstanceBuffer.mNumInstances;
    ID3D11Buffer& instanceBuffer = *Globals::gGeomBuffers.mInstanceBuffer.mBuffer.mBuffer;

    FrameInfo* forRendering = Globals::gFrameBuffer.mForRendering;
    copyData(forRendering->mEntitiesData, sizeof(KinematicInfo) * numInstances, instanceBuffer);
}